// Load debug symbols to HIEW as local names

// ----------------------------------------------------------------------------

#include "stdafx.h"

#include "DSymLoadEnum.h"
#include "DSymLoadError.h"

#include "DSymLoadImpl.h"

// ----------------------------------------------------------------------------

namespace DSymLoad
{

// ----------------------------------------------------------------------------

namespace
{

// ----------------------------------------------------------------------------

struct CSymbolNameLess
{
    bool operator()(const CStringA &_Left,const CStringA &_Right) const
    {
        return _Left.CompareNoCase(_Right) < 0;
    }
};

// ----------------------------------------------------------------------------

class CLocalNamesHarvester : public IEnumSymbols
{
public:
    explicit CLocalNamesHarvester(const CSettings &Settings);

    virtual void Callback(__in PSYMBOL_INFO pSymInfo, __in ULONG SymbolSize) override;

    typedef std::map< CStringA, ULONG64, CSymbolNameLess > CCollection;
    const CCollection &GetCollection() const;

private:
    const CSettings &m_Settings;

    CCollection m_Collection;
};

// ----------------------------------------------------------------------------

CLocalNamesHarvester::CLocalNamesHarvester(const CSettings &Settings)
    : m_Settings(Settings)
{
}

// ----------------------------------------------------------------------------

void CLocalNamesHarvester::Callback(__in PSYMBOL_INFO pSymInfo, __in ULONG SymbolSize)
{
    UNREFERENCED_PARAMETER(SymbolSize);
    if (!pSymInfo->NameLen || !*pSymInfo->Name || !pSymInfo->Address)
        return;

    if (pSymInfo->Flags & SYMFLAG_TLSREL)
        return;

    CStringA csName = &pSymInfo->Name[0];

    const auto it = m_Collection.find(csName);
    if (it == m_Collection.end())
    {
        m_Collection[csName] = pSymInfo->Address;
        return;
    }
    if (!m_Settings.IsProcessNameCollision())
        return;

    csName.AppendFormat("_%I64X", pSymInfo->Address);
    m_Collection[csName] = pSymInfo->Address;
}

// ----------------------------------------------------------------------------

const CLocalNamesHarvester::CCollection &CLocalNamesHarvester::GetCollection() const
{
    return m_Collection;
}

// ----------------------------------------------------------------------------

}

// ----------------------------------------------------------------------------

void Do(__in const CSettings &Settings, __in PCSTR szFileName)
{
    CLocalNamesHarvester Harvester(Settings);
    EnumAllSymbols(Settings, szFileName, Harvester);

    for (const auto &Name : Harvester.GetCollection())
    {
        std::vector<HEM_BYTE> szNameBuffer(Name.first.GetLength() + 1, '\0');
        std::copy(
            Name.first.GetString(),
            Name.first.GetString() + Name.first.GetLength(),
            szNameBuffer.begin());

        const auto Result = HiewGate_Names_AddLocal(Name.second, &szNameBuffer[0]);
        if ((HEM_ERROR == Result) && Settings.IsOverrideIfExist())
        {
            int bLocal = TRUE;
            const auto pFounded = 
                HiewGate_Names_FindName(
                    &szNameBuffer[0],
                    &bLocal);
            if (HEM_OFFSET_NOT_FOUND != pFounded && bLocal)
            {
                HiewGate_Names_DelLocal(pFounded);
                HiewGate_Names_AddLocal(Name.second, &szNameBuffer[0]);
            }
        }
    }
}

// ----------------------------------------------------------------------------

}   // namespace DSymLoad

// ----------------------------------------------------------------------------
